/* Copyright (C)
 * 2020 - ConleyLee
 * This program is free software; you can redistribute it and/or
 * modify it under the terms of the GNU General Public License
 * as published by the Free Software Foundation; either version 2
 * of the License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.
 *
 */

#pragma once
#include "Eigen/Dense"
#include "ps/ps.h"
#include <assert.h>
#include <leveldb/db.h>
//#include <glog/logging.h>

namespace ps {
namespace tensorflow {

    /**
     * @brief 用于支持tensorflow的SGD优化器，tensorflow worker推送梯度到SGDServer之后，
     * 更新参数，并pull最新参数到worker端
     *
     * \code
     *  auto server = std::make_shared<SGDServer<float>>(app_id, "./parameters.db", 0.001)
     * \endcode
     * @tparam Val 参数的类型
     */
template<class Val>
class SGDServer
{
  public:
    using Matrix =
      Eigen::Matrix<Val, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>;
    SGDServer(int app_id, std::string const& db, double const lr = 0.01)
      : _lr(lr)
    {

        leveldb::Options opt;
        opt.create_if_missing = true;
        leveldb::Status status = leveldb::DB::Open(opt, db, &_db);
        assert(status.ok());
        // LOG(INFO) << "Open parameters database at: " << db;

        std::string value;
        status = _db->Get(leveldb::ReadOptions(), "global_steps", &value);
        assert(status.ok() || status.IsNotFound());
        if (status.IsNotFound()) {
            _global_steps = 0;
            leveldb::Slice v((char*)(&_global_steps), sizeof(_global_steps));
            _db->Put(leveldb::WriteOptions(), "global_steps", v);
        } else {
            _global_steps = *(uint64_t*)value.data();
        }

        ps::KVServer<Val>* server = new ps::KVServer<Val>(app_id);
        using namespace std::placeholders;
        server->set_request_handle(
          std::bind(&SGDServer::req_handler, this, _1, _2, _3));
        ps::RegisterExitCallback([server, this]() {
            delete server;
            delete _db;
        });
    }

  private:
    double const _lr;
    leveldb::DB* _db;
    uint64_t _global_steps;

    /**
     * @brief handle the request from workers.
     *
     * @param meta meta information for ps-lite
     * @param The data pushed by workers
     * @param server server instance
     */
    void req_handler(ps::KVMeta const& meta,
                     ps::KVPairs<Val> const& data,
                     ps::KVServer<Val>* server)
    {
        ps::KVPairs<Val> res;
        if (meta.push) {
            LOG(INFO) << "receive push";
            do_push(data);
        }

        if (meta.pull) {
            LOG(INFO) << "receive pull";
            do_pull(data, res);
        }

        server->Response(meta, res);
    }

    /**
     * @brief update parameters
     *
     * @param kvs parameters post by client.
     */
    void do_push(ps::KVPairs<Val> const& kvs)
    {
        Val const* pos = kvs.vals.begin();
        for (auto i = 0; i < kvs.keys.size(); ++i) {
            auto k = kvs.keys[i];
            auto n = kvs.lens[i];

            std::string value;
            std::string const key(std::to_string(k));

            leveldb::Status s = _db->Get(leveldb::ReadOptions(), key, &value);

            if (s.IsNotFound()) {
                leveldb::Slice v((char*)pos, n * sizeof(Val));
                _db->Put(leveldb::WriteOptions(), key, v);
            } else {
                assert(value.size() / sizeof(Val) == n);
                Eigen::Map<const Matrix> gradient(pos, 1, n);

                // std::string is required to use contiguous storage in c++11
                Eigen::Map<Matrix> var((Val*)(value.data()), 1, n);
                var -= _lr * gradient; //直接增加
                s = _db->Put(leveldb::WriteOptions(),
                             key,
                             leveldb::Slice(value.data(), value.size()));
                assert(s.ok());
            }
        }
    }

    /**
     * @brief
     *
     * @param kvs
     */
    void do_pull(ps::KVPairs<Val> const& kvs, ps::KVPairs<Val>& res)
    {
        std::vector<Val> vals;

        for (auto i = 0; i < kvs.keys.size(); ++i) {
            auto k = kvs.keys[i];
            // LOG(INFO) << "pull key " << k;
            std::string const key(std::to_string(k));
            std::string value;
            auto status = _db->Get(leveldb::ReadOptions(), key, &value);

            if (status.IsNotFound()) {
                LOG(INFO) << "key " << k << " not found.";
                res.keys.push_back(k);
                res.lens.push_back(0);
                continue;
            }

            auto n = value.size() / sizeof(Val);
            LOG(INFO) << "get key " << key << ", size = " << n;
            res.keys.push_back(k);
            res.lens.push_back(n);
            std::copy((Val*)(value.data()),
                      (Val*)(value.data()) + n,
                      std::back_inserter(vals));
        }

        LOG(INFO) << "after process, vals.size() = " << vals.size();
        res.vals.CopyFrom(vals.cbegin(), vals.cend());
    }
};
}
}

